import torch
import matplotlib.pyplot as plt

"""
优化器：用于更新模型的参数以最小化损失函数，提高模型的准确性和性能。
GD算法：一种基本的一阶优化算法。w新 = w旧 - lr * 梯度
"""


# 定义损失函数
def loss_fn(w1, w2):
    return w1 ** 2 + w2 ** 2


# 超参数
lr = 0.1
Epochs = 20
w1 = -1
w2 = 1

'''绘制等高线图'''
x1 = torch.linspace(-1, 1, 100)
x2 = torch.linspace(-1, 1, 100)
xx1, xx2 = torch.meshgrid(x1, x2, indexing='ij')

fig = plt.figure(figsize=(12, 6))
ax = fig.add_subplot()
ax.contourf(xx1, xx2, loss_fn(xx1, xx2), cmap='rainbow')

points = []
# 循环训练
for epoch in range(Epochs):
    points.append([w1, w2])
    loss = loss_fn(w1, w2)
    print(loss)
    # gd算法
    g1 = 2 * w1
    g2 = 4 * w2
    w1 = w1 - lr * g1
    w2 = w2 - lr * g2

points = torch.tensor(points)
ax.plot(points[:, 0], points[:, 1], 'o-', color='r')
plt.show()